{-# LANGUAGE TemplateHaskell #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Wire.DomainRegistrationStore
  ( DomainRegistrationStore (..),
    StoredDomainRegistration (..),
    DomainKey,
    upsert,
    lookup,
    lookupByTeam,
    delete,
  )
where

import Cassandra
import Data.ByteString.Conversion
import Data.CaseInsensitive
import Data.CaseInsensitive qualified as CI
import Data.Domain as Domain
import Data.Id
import Data.Misc
import Data.Text as T
import Database.CQL.Protocol (Record (..), TupleType, recordInstance)
import Imports hiding (lookup)
import Polysemy
import Polysemy.Error
import Polysemy.TinyLog (TinyLog)
import Polysemy.TinyLog qualified as Log
import SAML2.WebSSO qualified as SAML
import System.Logger.Message qualified as Log
import Wire.API.EnterpriseLogin

newtype DomainKey = DomainKey {unDomainKey :: CI Text}
  deriving stock (Eq, Ord, Show)

mkDomainKey :: Domain -> DomainKey
mkDomainKey = DomainKey . CI.mk . domainText

unmkDomainKey :: DomainKey -> Domain
unmkDomainKey = Domain . CI.foldedCase . unDomainKey

instance Cql DomainKey where
  ctype = Tagged TextColumn
  toCql = CqlText . CI.foldedCase . unDomainKey
  fromCql (CqlText txt) = pure . DomainKey . CI.mk $ txt
  fromCql _ = Left "DomainKey: Text expected"

data StoredDomainRegistration = StoredDomainRegistration
  { domain :: DomainKey,
    domainRedirect :: Maybe DomainRedirectTag,
    teamInvite :: Maybe TeamInviteTag,
    idpId :: Maybe SAML.IdPId,
    backendUrl :: Maybe HttpsUrl,
    team :: Maybe TeamId,
    dnsVerificationToken :: Maybe DnsVerificationToken,
    authTokenHash :: Maybe Token,
    authorizedTeam :: Maybe TeamId,
    webappUrl :: Maybe HttpsUrl
  }
  deriving (Show, Eq, Ord, Generic)

recordInstance ''StoredDomainRegistration

data DomainRegistrationStore m a where
  UpsertInternal :: StoredDomainRegistration -> DomainRegistrationStore m ()
  LookupInternal :: DomainKey -> DomainRegistrationStore m (Maybe StoredDomainRegistration)
  LookupByTeamInternal :: TeamId -> DomainRegistrationStore m [StoredDomainRegistration]
  DeleteInternal :: DomainKey -> DomainRegistrationStore m ()

upsert :: (Member DomainRegistrationStore r) => DomainRegistration -> Sem r ()
upsert = send . UpsertInternal . toStored

lookupByTeam :: forall r. (Member DomainRegistrationStore r, Member (Log.Logger (Log.Msg -> Log.Msg)) r) => TeamId -> Sem r [DomainRegistration]
lookupByTeam tid = do
  rows <- send (LookupByTeamInternal tid)
  mRegisteredDomains <- for rows fromStoredWithLogging
  pure $ catMaybes mRegisteredDomains
  where
    fromStoredWithLogging :: StoredDomainRegistration -> Sem r (Maybe DomainRegistration)
    fromStoredWithLogging row = case fromStored row of
      Just dr -> pure (Just dr)
      Nothing -> logInvalidDomainRegistrationError (unmkDomainKey row.domain) $> Nothing

lookup ::
  forall r.
  (Member DomainRegistrationStore r, Member TinyLog r) =>
  Domain ->
  Sem r (Maybe DomainRegistration)
lookup domain =
  (>>= logErrors) . runError @Bool $ do
    sdr <- send (LookupInternal (mkDomainKey domain)) >>= note False
    fromStored sdr & note True
  where
    logErrors :: Either Bool a -> Sem r (Maybe a)
    logErrors (Left False) = pure Nothing
    logErrors (Left True) = logInvalidDomainRegistrationError domain $> Nothing
    logErrors (Right x) = pure (Just x)

logInvalidDomainRegistrationError :: (Member TinyLog r, ToByteString a) => a -> Sem r ()
logInvalidDomainRegistrationError domain =
  Log.err $
    Log.field "domain" (toByteString' domain)
      . Log.msg (Log.val "Invalid stored domain registration")

delete :: (Member DomainRegistrationStore r) => Domain -> Sem r ()
delete = send . DeleteInternal . mkDomainKey

fromStored :: StoredDomainRegistration -> Maybe DomainRegistration
fromStored sdr =
  DomainRegistration (unmkDomainKey sdr.domain) sdr.authorizedTeam
    <$> getDomainRedirect sdr
    <*> getTeamInvite sdr
    <*> pure sdr.dnsVerificationToken
    <*> pure sdr.authTokenHash
  where
    getTeamInvite :: StoredDomainRegistration -> Maybe TeamInvite
    getTeamInvite = \case
      StoredDomainRegistration _ _ ti _ _ tid _ _ _ _ -> case (ti, tid) of
        (Just AllowedTag, Nothing) -> Just Allowed
        (Just NotAllowedTag, Nothing) -> Just NotAllowed
        (Just TeamTag, Just teamId) -> Just $ Team teamId
        (Nothing, Nothing) -> Just Allowed
        _ -> Nothing

    getDomainRedirect :: StoredDomainRegistration -> Maybe DomainRedirect
    getDomainRedirect = \case
      StoredDomainRegistration _ dr _ ssoId url _ _ _ _ webappUrl -> case (dr, ssoId, url, webappUrl) of
        (Just NoneTag, Nothing, Nothing, Nothing) -> Just None
        (Just LockedTag, Nothing, Nothing, Nothing) -> Just Locked
        (Just PreAuthorizedTag, Nothing, Nothing, Nothing) -> Just PreAuthorized
        (Just SSOTag, Just idpId, Nothing, Nothing) -> Just $ SSO idpId
        (Just BackendTag, Nothing, Just beUrl, webAppUrl) -> Just $ Backend beUrl webAppUrl
        (Just NoRegistrationTag, Nothing, Nothing, Nothing) -> Just NoRegistration
        (Nothing, Nothing, Nothing, Nothing) -> Just None
        _ -> Nothing

toStored :: DomainRegistration -> StoredDomainRegistration
toStored dr =
  let (domainRedirect, idpId, backendUrl, webappUrl) = fromDomainRedirect dr.domainRedirect
      (teamInvite, team) = fromTeamInvite dr.teamInvite
   in StoredDomainRegistration
        { domain = mkDomainKey dr.domain,
          domainRedirect = Just domainRedirect,
          teamInvite = Just teamInvite,
          idpId,
          backendUrl,
          team,
          dnsVerificationToken = dr.dnsVerificationToken,
          authTokenHash = dr.authTokenHash,
          authorizedTeam = dr.authorizedTeam,
          webappUrl
        }
  where
    fromTeamInvite :: TeamInvite -> (TeamInviteTag, Maybe TeamId)
    fromTeamInvite Allowed = (AllowedTag, Nothing)
    fromTeamInvite NotAllowed = (NotAllowedTag, Nothing)
    fromTeamInvite (Team teamId) = (TeamTag, Just teamId)

    fromDomainRedirect :: DomainRedirect -> (DomainRedirectTag, Maybe SAML.IdPId, Maybe HttpsUrl, Maybe HttpsUrl)
    fromDomainRedirect None = (NoneTag, Nothing, Nothing, Nothing)
    fromDomainRedirect Locked = (LockedTag, Nothing, Nothing, Nothing)
    fromDomainRedirect (SSO idpId) = (SSOTag, Just idpId, Nothing, Nothing)
    fromDomainRedirect (Backend url webappUrl) = (BackendTag, Nothing, Just url, webappUrl)
    fromDomainRedirect NoRegistration = (NoRegistrationTag, Nothing, Nothing, Nothing)
    fromDomainRedirect PreAuthorized = (PreAuthorizedTag, Nothing, Nothing, Nothing)
